package data_structure.mooc;

import java.util.TreeMap;

/**
 * 字典树
 */
public class Trie {

    /**
     * 节点类
     */
    private class Node{
        //记录从根节点到当前访问的位置中的元素是否为单词
        public boolean isWord;
        //记录下一个节点
        public TreeMap<Character,Node>next;

        public Node(boolean isWord){
            this.isWord = isWord;
            next=new TreeMap<Character, Node> ();
        }
        public Node(){
            this(false);
        }
    }

    /**
     * 字典树成员变量
     */
    private Node root;
    private int size;

    public Trie(Node root){
        this.root = root;
        size=0;
    }

    public Trie(){
        this(null);
    }

    public int getSize(){
        return size;
    }

    /**
     * 向字典树中添加字符串
     * @param string
     */
    public void add(String string){
        Node curr=root;
        for(int i=0;i<string.length ();i++){
            char c=string.charAt(i);
            if(curr.next.get(c)==null){
                curr.next.put (c,new Node());
            }
            curr=curr.next.get (c);
        }
        if(!curr.isWord){
            curr.isWord=true;
            size++;
        }
    }

    /**
     * 递归实现向树中添加元素
     */
    public void addR(String string){
        Node curr=root;
        addR(curr,string,0);
    }

    private void addR(Node root,String string,int index){

        if(index==string.length()&&!root.isWord){
            root.isWord=true;
            size++;
            return;
        }
        if(string.length ()>index){

            char c=string.charAt(index);
            if(root.next.get (c)==null){
                root.next.put(c,new Node());
            }
            addR (root.next.get(c),string,index+1);
        }
    }

    /**
     * 查找单词
     * 从根节点开始查找，将当前字符作为key值，跳转到指定的字节点，如果子节点不存在，返回false
     */
    public boolean contains(String word){
        Node curr=root;
        for(int i=0;i<word.length(); i++){
            char c=word.charAt(i);
            if(root.next.get(c) == null)return false;
            curr=curr.next.get(c);
        }
        //遍历结束后，判断从root到curr是否为一个单词
        return curr.isWord;
    }

    /**
     * 查询字典树中是否包含前缀prefix，和contains方法逻辑处理相同
     */
    public boolean isPrefix(String prefix){
        Node curr=root;
        for(int i = 0; i < prefix.length(); i++){
            char c=prefix.charAt(i);
            if(root.next.get(c) == null)return false;
            curr = curr.next.get(c);
        }
        return true;
    }
}
